// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
//
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include <tt_stl/assert.hpp>
#include <tt-metalium/mesh_coord.hpp>
#include <tt-metalium/experimental/fabric/fabric_types.hpp>
#include <tt_stl/reflection.hpp>
#include <tt_stl/indestructible.hpp>
#include <umd/device/types/arch.hpp>                      // tt::ARCH
#include <umd/device/types/cluster_descriptor_types.hpp>  // ChipId
#include <cstddef>
#include <cstdint>
#include <functional>
#include <string>
#include <unordered_map>
#include <utility>

#include <tt-metalium/experimental/fabric/mesh_graph_descriptor.hpp>

#include <vector>

namespace tt {
namespace tt_metal {
enum class ClusterType : std::uint8_t;
}  // namespace tt_metal
}  // namespace tt
namespace tt::tt_fabric {

using tt::tt_metal::distributed::MeshContainer;
using tt::tt_metal::distributed::MeshCoordinate;
using tt::tt_metal::distributed::MeshCoordinateRange;
using tt::tt_metal::distributed::MeshShape;

struct ChipSpec {
    tt::ARCH arch;
    std::uint32_t num_eth_ports_per_direction;
    std::uint32_t num_z_ports;
};

enum class RoutingDirection {
    N = 0,
    E = 1,
    S = 2,
    W = 3,
    Z = 4,
    C = 5,     // Centre, means that destination is same as source
    NONE = 6,  // No direction, means that destination is not reachable
};

struct RouterEdge {
    // TODO: change this to be port_id_t
    RoutingDirection port_direction;  // Assume all ports in one direction connect to the same chip
    std::vector<ChipId>
        connected_chip_ids;  // One per port to the connected chip, used by ControlPlane to map to physical links
    std::uint32_t weight;    // Assume all chip to chip communication is equal weight, but in reality there may be less
                             // intermesh traffic for example
};
struct hash_pair {
    template <class T1, class T2>
    size_t operator()(const std::pair<T1, T2>& p) const {
        return tt::stl::hash::hash_objects(std::hash<T1>{}(p.first), std::hash<T2>{}(p.second));
    }
};

using port_id_t = std::pair<RoutingDirection, uint32_t>;

}  // namespace tt::tt_fabric

// Custom hash function for port_id_t to use in unordered containers
namespace std {

template <>
struct hash<tt::tt_fabric::port_id_t> {
    std::size_t operator()(const tt::tt_fabric::port_id_t& p) const;
};

}  // namespace std

namespace tt::tt_fabric {
using InterMeshConnectivity = std::vector<std::vector<std::unordered_map<MeshId, RouterEdge>>>;
using IntraMeshConnectivity = std::vector<std::vector<std::unordered_map<ChipId, RouterEdge>>>;
// Vector of Intermesh Connections generated by the Control Plane based on high level user input (MGD).
// Each connection is fully specified by the meshes it connects and the logical ports (direction and logical channels)
// it uses.
using AnnotatedIntermeshConnections =
    std::vector<std::tuple<std::pair<uint32_t, port_id_t>, std::pair<uint32_t, port_id_t>>>;

// Parsed from the Mesh Graph Descriptor. The user can specify the number of channels betweeen meshes (relaxed mode)
// or pin connectiosn to specific exit nodes (strict mode).
// Stores connections specified in relaxed mode. Mapping: src_mesh -> dst_mesh -> num_channels
using RequestedIntermeshConnections = std::unordered_map<uint32_t, std::unordered_map<uint32_t, uint32_t>>;

// Stores connections specified in strict mode. Mapping: src_mesh -> dst_mesh -> list of (src_device, dst_device,
// num_channels)
using RequestedIntermeshPorts =
    std::unordered_map<uint32_t, std::unordered_map<uint32_t, std::vector<std::tuple<uint32_t, uint32_t, uint32_t>>>>;

class MeshGraph {
public:
    explicit MeshGraph(
        const std::string& mesh_graph_desc_file_path, std::optional<FabricConfig> fabric_config = std::nullopt);
    MeshGraph() = delete;
    ~MeshGraph() = default;

    void print_connectivity() const;

    const IntraMeshConnectivity& get_intra_mesh_connectivity() const;
    const InterMeshConnectivity& get_inter_mesh_connectivity() const;

    const ChipSpec& get_chip_spec() const { return chip_spec_; }

    // Get the host ranks for a given mesh_id
    // Returned MeshContainer has a shape denoting the shape of how the "board" are arranged
    const MeshContainer<MeshHostRankId>& get_host_ranks(MeshId mesh_id) const;

    // Get the shape of the mesh, or the shape of the submesh for a given host rank if provided
    MeshShape get_mesh_shape(MeshId mesh_id, std::optional<MeshHostRankId> host_rank = std::nullopt) const;

    // Get the coordinate range of the mesh, or the coordinate range of the submesh for a given host rank if provided
    MeshCoordinateRange get_coord_range(MeshId mesh_id, std::optional<MeshHostRankId> host_rank = std::nullopt) const;

    std::vector<MeshId> get_mesh_ids() const;

    // Get the chip ids for a given mesh_id
    // If host_rank is provided, return the chip ids for the submesh for that host rank
    // Otherwise, return the chip ids for the entire mesh
    MeshContainer<ChipId> get_chip_ids(MeshId mesh_id, std::optional<MeshHostRankId> host_rank = std::nullopt) const;

    // Switch query APIs (internal representation uses MeshId, API uses SwitchId)
    std::vector<SwitchId> get_switch_ids() const;
    std::unordered_set<MeshId> get_meshes_connected_to_switch(SwitchId switch_id) const;
    bool is_mesh_connected_to_switch(MeshId mesh_id, SwitchId switch_id) const;
    std::optional<SwitchId> get_switch_for_mesh(MeshId mesh_id) const;

    // Get the host rank that owns a given chip in a mesh
    std::optional<MeshHostRankId> get_host_rank_for_chip(MeshId mesh_id, ChipId chip_id) const;

    // Translation functions for chip_id and coordinate using RM-convention
    MeshCoordinate chip_to_coordinate(MeshId mesh_id, ChipId chip_id) const;
    ChipId coordinate_to_chip(MeshId mesh_id, MeshCoordinate coordinate) const;

    // Static functions for mesh graph descriptor management
    static std::filesystem::path get_mesh_graph_descriptor_path_for_cluster_type(
        tt::tt_metal::ClusterType cluster_type,
        const std::string& root_dir,
        tt::tt_fabric::FabricType fabric_type = tt::tt_fabric::FabricType::MESH);

    // Get the number of active channels the user has requested between meshes
    const RequestedIntermeshConnections& get_requested_intermesh_connections() const;

    // Get the number of active channels the user has requested between specific logical devices across meshes
    const RequestedIntermeshPorts& get_requested_intermesh_ports() const;

    // Query the mapping of logical ports to logical device ids per mesh
    const std::vector<std::unordered_map<port_id_t, ChipId, hash_pair>>& get_mesh_edge_ports_to_chip_id() const;

    // Load Inter-Mesh Connectivity into the Mesh Graph.
    void load_intermesh_connections(const AnnotatedIntermeshConnections& intermesh_connections);

    bool is_intra_mesh_policy_relaxed(MeshId mesh_id) const;

private:
    void validate_mesh_id(MeshId mesh_id) const;
    std::unordered_map<ChipId, RouterEdge> get_valid_connections(
        const MeshCoordinate& src_mesh_coord,
        const MeshCoordinateRange& mesh_coord_range,
        FabricType fabric_type) const;
    void initialize_from_mgd(const MeshGraphDescriptor& mgd, std::optional<FabricConfig> fabric_config);

    void add_to_connectivity(
        MeshId src_mesh_id,
        ChipId src_chip_id,
        MeshId dest_mesh_id,
        ChipId dest_chip_id,
        RoutingDirection port_direction);

    ChipSpec chip_spec_{};
    std::map<MeshId, MeshContainer<ChipId>> mesh_to_chip_ids_;
    IntraMeshConnectivity intra_mesh_connectivity_;
    InterMeshConnectivity inter_mesh_connectivity_;

    // For distributed context, bookkeeping of host ranks and their shapes
    std::vector<MeshContainer<MeshHostRankId>> mesh_host_ranks_;
    std::unordered_map<std::pair<MeshId, MeshHostRankId>, MeshCoordinateRange, hash_pair> mesh_host_rank_coord_ranges_;

    std::vector<std::unordered_map<port_id_t, ChipId, hash_pair>> mesh_edge_ports_to_chip_id_;
    RequestedIntermeshConnections requested_intermesh_connections_;
    RequestedIntermeshPorts requested_intermesh_ports_;

    // Switch tracking (switches use MeshId as their identifier)
    std::vector<MeshId> switch_ids_;
    std::map<MeshId, MeshContainer<ChipId>> switch_to_chip_ids_;
    std::unordered_map<MeshId, std::vector<MeshId>> switch_to_connected_meshes_;
    std::unordered_map<MeshId, bool> intra_mesh_relaxed_policy_;
};

}  // namespace tt::tt_fabric
